"""
Reinforcement learning maze example.

Red rectangle:          explorer.
Black rectangles:       hells       [reward = -1].
Yellow bin circle:      paradise    [reward = +1].
All other states:       ground      [reward = 0].

This script is the environment part of this example. The RL is in RL_brain.py.

View more on my tutorial page: https://morvanzhou.github.io/tutorials/
"""


import numpy as np
import time
import sys
if sys.version_info.major == 2:
    import Tkinter as tk
else:
    import tkinter as tk


UNIT = 40   # pixels
MAZE_H = 9  # grid height 不大于10
MAZE_W = 9  # grid width 不大于10


class Maze(tk.Tk, object):
    def __init__(self):
        super(Maze, self).__init__()
        self.action_space = ['u', 'd', 'l', 'r']
        self.n_actions = len(self.action_space)
        self.title('maze')
        self.geometry('{0}x{1}'.format(MAZE_H * UNIT, MAZE_H * UNIT))

        self.hell = [[0, 3], [3, 0], [3, 3], [6, 1], [1, 6], [4, 6], [8, 3], [3, 8],
                     [6, 4], [7, 6], [5, 2], [2, 5], [6, 7], [1, 1]] # 黑洞位置
        self.oval1 = [7, 7] # 目标位置
        self.rect1 = [0, 0] # 初始位置

        self._build_maze()


    def _build_maze(self):
        self.canvas = tk.Canvas(self, bg='white',
                           height=MAZE_H * UNIT,
                           width=MAZE_W * UNIT)

        for c in range(0, MAZE_W * UNIT, UNIT): # create grids
            x0, y0, x1, y1 = c, 0, c, MAZE_H * UNIT
            self.canvas.create_line(x0, y0, x1, y1)
        for r in range(0, MAZE_H * UNIT, UNIT):
            x0, y0, x1, y1 = 0, r, MAZE_W * UNIT, r
            self.canvas.create_line(x0, y0, x1, y1)

        for _i in range(len(self.hell)): # 创建黑洞
            self.canvas.create_rectangle(
                self.hell[_i][1]*UNIT + 5, self.hell[_i][0]*UNIT + 5,
                self.hell[_i][1]*UNIT + 35, self.hell[_i][0]*UNIT + 35,
                fill='black')
            
        self.canvas.create_oval( # 创建标
            self.oval1[1]*UNIT + 5, self.oval1[0]*UNIT + 5,
            self.oval1[1]*UNIT + 35, self.oval1[0]*UNIT + 35,
            fill='yellow')

        self.rect = self.canvas.create_rectangle(  # 创建主体
            self.rect1[1]*UNIT + 5, self.rect1[0]*UNIT + 5,
            self.rect1[1]*UNIT + 35, self.rect1[0]*UNIT + 35,
            fill='red')

        self.canvas.pack() # pack all


    def reset(self, rand): # 每个回合初始状态设置
        time.sleep(0.2)
        if rand: # 设置随机初始值
            while True:
                a = np.random.randint(MAZE_W * MAZE_H)
                self.rect1 = [int(a/MAZE_H), int(a%MAZE_H)] # 随机取初始值
                if (self.rect1 not in self.hell) and (self.rect1 != self.oval1):
                    break;
        else: # 恢复初始值
            self.rect1 = [0, 0]

        self.canvas.delete(self.rect) # 删去上个回合主体
        self.rect = self.canvas.create_rectangle( # 重新创建主体
            self.rect1[1]*UNIT + 5, self.rect1[0]*UNIT + 5,
            self.rect1[1]*UNIT + 35, self.rect1[0]*UNIT + 35,
            fill='red')
        
        self.update()
        return self.rect1


    def step(self, action): # 行动
        s = self.rect1.copy()
        if action == 0:   # up
            if s[0] > 0:
                s[0] -= 1
        elif action == 1:   # down
            if s[0]  < (MAZE_H - 1):
                s[0] += 1
        elif action == 2:   # right
            if s[1]  < (MAZE_W - 1):
                s[1] += 1
        elif action == 3:   # left
            if s[1]  >0:
                s[1] -= 1

        if (s != self.rect1): # 若有变动，更新状态
            self.canvas.move(self.rect,
                    (s[0]-self.rect1[0])*UNIT, (s[1]-self.rect1[1])*UNIT)  # move agent
            self.rect1 = s.copy()
            reward = 0
        else:
            reward = -1 # 调整到前面，便于边界越线扣分，提高前期探索效率

        done = False
        if self.rect1 == self.oval1: # 达到目标
            reward = 10
            done = True
        elif self.rect1 in self.hell: # 掉进黑洞
            reward = -10
            done = True

        self.update() # 调整到子程序最后更新
        return self.rect1, reward, done


    def render(self): # 更新界面
        time.sleep(0.01)
        self.update()


def update1():
    for t in range(10):
        s = env.reset()
        while True:
            env.render()
            a = np.random.randint(0,4)
            s, r, done = env.step(a)
            if done:
                break
    print('game over') 
    env.destroy()


if __name__ == '__main__':
    env = Maze()
    env.after(100, update1)
    env.mainloop()
